/// LM 算法

#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;

/// LM callback - 单应矩阵
/// opencv 4 版本以后 cv::findHomography 就是使用的 这个方法计算
class HomographyRefineCallback CV_FINAL : public LMSolver::Callback
{
public:
    HomographyRefineCallback(InputArray _src, InputArray _dst)
    {
        src = _src.getMat();
        dst = _dst.getMat();
    }

    bool compute(InputArray _param, OutputArray _err, OutputArray _Jac) const CV_OVERRIDE
    {
        int i, count = src.checkVector(2);
        Mat param = _param.getMat();
        _err.create(count*2, 1, CV_64F);
        Mat err = _err.getMat(), J;
        if( _Jac.needed())
        {
            _Jac.create(count*2, param.rows, CV_64F);
            J = _Jac.getMat();
            CV_Assert( J.isContinuous() && J.cols == 8 );
        }

        const Point2f* M = src.ptr<Point2f>();
        const Point2f* m = dst.ptr<Point2f>();
        const double* h = param.ptr<double>();
        double* errptr = err.ptr<double>();
        double* Jptr = J.data ? J.ptr<double>() : 0;

        for( i = 0; i < count; i++ )
        {
            double Mx = M[i].x, My = M[i].y;
            double ww = h[6]*Mx + h[7]*My + 1.;
            ww = fabs(ww) > DBL_EPSILON ? 1./ww : 0;
            double xi = (h[0]*Mx + h[1]*My + h[2])*ww;
            double yi = (h[3]*Mx + h[4]*My + h[5])*ww;
            errptr[i*2] = xi - m[i].x;
            errptr[i*2+1] = yi - m[i].y;

            if( Jptr )
            {
                Jptr[0] = Mx*ww; Jptr[1] = My*ww; Jptr[2] = ww;
                Jptr[3] = Jptr[4] = Jptr[5] = 0.;
                Jptr[6] = -Mx*ww*xi; Jptr[7] = -My*ww*xi;
                Jptr[8] = Jptr[9] = Jptr[10] = 0.;
                Jptr[11] = Mx*ww; Jptr[12] = My*ww; Jptr[13] = ww;
                Jptr[14] = -Mx*ww*yi; Jptr[15] = -My*ww*yi;

                Jptr += 16;
            }
        }

        return true;
    }

    Mat src, dst;
};

/// 对比 单应矩阵 计算 1. findHomography. 2. LMSolver
static void demoH()
{
    vector<Point2f> pts_src = {
        {83.3333, 83.3333},
        {83.3333, 416.667},
        {416.667, 416.667},
        {416.667, 83.3333}
    };

    vector<Point2f> pts_dst = {
        {153.801, 217.638},
        {170.746, 393.078},
        {321.615, 393.563},
        {340.058, 215.632}
    };

    Mat H= findHomography(pts_src, pts_dst, RANSAC);
    cout << "H:\n" << H << endl;

    Mat H8(8, 1, CV_64F, H.ptr<double>());
    Mat src, dst;
    src = Mat(pts_src);
    dst = Mat(pts_dst);
    LMSolver::create(makePtr<HomographyRefineCallback>(src, dst), 10)->run(H8);

    cout << "H8:\n" << H8 << endl;
}


/// LM callback  解方程 y=exp(ax^2+b+c) 的参数 [a, b, c]
class SolverAbcCallBack : public LMSolver::Callback
{
public:

    SolverAbcCallBack(InputArray _src, InputArray _dst)
    {
        src = _src.getMat();
        dst = _dst.getMat();
    }

    bool compute(InputArray _param, OutputArray _err, OutputArray _Jac) const CV_OVERRIDE
    {
        // y=exp(ax^2+b+c)
        int i, count = src.cols;
        Mat param = _param.getMat();
        cout << "param: " << param << endl;
        _err.create(count, 1, CV_64F);
        Mat err = _err.getMat(), J;
        if( _Jac.needed())
        {
            _Jac.create(count, param.rows, CV_64F);
            J = _Jac.getMat();
            CV_Assert( J.isContinuous()  );
        }

        const double* x = src.ptr<double>();
        const double* y = dst.ptr<double>();
        const double* h = param.ptr<double>();  // 参数
        double* errptr = err.ptr<double>();
        double* Jptr = J.data ? J.ptr<double>() : 0;

        for( i = 0; i < count; i++ )
        {
            double xi = exp (h[0]*x[i]*x[i] + h[1]*x[i] + h[2]);
            errptr[i] = y[i] - xi;

            // 雅格比矩阵计算
            if( Jptr )
            {
                Jptr[0] = -x[i]*x[i]*exp(h[0]*x[i]*x[i] + h[1]*x[i] + h[2]);
                Jptr[1] = -x[i]*exp(h[0]*x[i]*x[i] + h[1]*x[i] + h[2]);
                Jptr[2] = -exp(h[0]*x[i]*x[i] + h[1]*x[i] + h[2]);
                Jptr+=3;
            }
        }
        return true;
    }

    Mat src, dst;
};

/// 解方程 y=exp(ax^2+b+c) 的参数 [a, b, c]
static void demoAbc()
{
    const double aa = 0.1, bb = 0.5, cc = 2; // 实际方程的参数
    double a =0.0, b=0.0, c=0.0; // 初值

    /* 制造数据 */
    const size_t N = 100; //数据个数
    cv::RNG rng(cv::getTickCount());
    vector<double> vecX;
    vector<double> vecY;
    double paramIni[3] = {a, b, c};
    for( size_t i = 0; i < N; i ++)
    {
        /* 生产带有高斯噪声的数据　*/
        double x = rng.uniform(0.0, 1.0) ;
        double y = exp(aa*x*x + bb*x + cc)/* + rng.gaussian(0.05)*/;
        vecX.push_back(x);
        vecY.push_back(y);
        /* 添加到观测中　*/
        // lm.addObservation(x, y);
    }

    Mat params = Mat(3, 1, CV_64F, paramIni);
    /* 构造问题 */
    LMSolver::create(makePtr<SolverAbcCallBack>(vecX, vecY), 200, 1e-10)->run(params);
    cout << "params: " << params << endl;

    cout << "Param result\n";
    for (int i = 0; i<3;i++)
        cout << paramIni[i] <<endl;
//    cout << "Error\n";
//    for (int i = 0; i<x.size();i++)
//        cout<<y[i]-( 1/(paramIni[0] * x[i] * x[i] + 1) + sin(paramIni[1] * x[i] + paramIni[2]))<<endl;
}

int main(int argc, char *argv[])
{
    /// H - opencv 自己的单应计算 和 自己写的LMSolver 计算
    demoH();
    // 解方程 y=exp(ax^2+b+c) 的参数 [a, b, c], 对比 [cvLMabc] https://zhuanlan.zhihu.com/p/42415718
    // demoAbc();
    return 0;
}
